from config.constant import LEAKAGE_MODEL, HW_ARRAY
from crypto_kit.cython.SM4_accelerator import SM4_SBOX, sm4_key_expansion, sm4_sbox_in
from crypto_kit.algorithm import CryptographyAlgorithm
import numpy as np

class SM4(CryptographyAlgorithm):
    @staticmethod
    def _func_real_enc_SboxInID(self):
        rkey_byte_need = int(self.target_byte / 4) * 4+4
        if self.rkey is None and rkey_byte_need > 0:
            raise RuntimeError(f'缺少轮密钥，请提供至少{rkey_byte_need}字节的轮密钥！')
        if rkey_byte_need > self.rkey.shape[1]:
            raise RuntimeError(f'轮密钥字节不足，请提供至少{rkey_byte_need}字节的轮密钥！')
        self.sbox_in = sm4_sbox_in(self.plaintext, self.target_byte, self.rkey)
        return self.sbox_in

    @staticmethod
    def _func_get_real_b0XORb02(self):
        B0=SM4_SBOX[self.plaintext[:,12]^self.rkey[:,0]]
        B02=(B0^(B0<<2))&0xff
        xor=self.plaintext[:,12]^B02
        return xor

    func_real_enc_SboxInID=lambda self:SM4._func_real_enc_SboxInID(self)
    func_real_enc_SboxInHW = lambda self: HW_ARRAY[SM4.func_real_enc_SboxInID(self)]
    func_real_enc_SboxInMBit = lambda self: np.array([[bit for bit in np.binary_repr(sbox_in, 8)] for sbox_in in SM4.func_real_enc_SboxInID(self)], np.uint8)
    func_real_enc_SboxOutID = lambda self: SM4_SBOX[SM4.func_real_enc_SboxInID(self)]
    func_real_enc_SboxOutHW = lambda self: HW_ARRAY[SM4.func_real_enc_SboxOutID(self)]
    func_real_enc_SboxOutMBit = lambda self: np.array([[bit for bit in np.binary_repr(sbox_out, 8)] for sbox_out in SM4.func_real_enc_SboxOutID(self)], np.uint8)
    func_real_enc_RoundXor=lambda self:SM4._func_get_real_b0XORb02(self)
    func_real_enc_RoundHD=lambda self:HW_ARRAY[SM4._func_get_real_b0XORb02(self)]
    func_real_enc_RoundMBit=lambda self: np.array([[bit for bit in np.binary_repr(xor, 8)] for xor in SM4._func_get_real_b0XORb02(self)], np.uint8)

    @staticmethod
    def _func_real_dec_SboxInID(self):
        rkey_byte_need = int(self.target_byte / 4) * 4 + 4
        if self.rkey is None and rkey_byte_need > 0:
            raise RuntimeError(f'缺少轮密钥，请提供至少{rkey_byte_need}字节的轮密钥！')
        if rkey_byte_need > self.rkey.shape[1]:
            raise RuntimeError(f'轮密钥字节不足，请提供至少{rkey_byte_need}字节的轮密钥！')
        return sm4_sbox_in(self.ciphertext, self.rkey, self.target_byte)

    func_real_dec_SboxInID=lambda self:SM4._func_real_dec_SboxInID(self)
    func_real_dec_SboxInHW = lambda self: HW_ARRAY[SM4.func_real_dec_SboxInID(self)]
    func_real_dec_SboxInMBit = lambda self: np.array([[bit for bit in np.binary_repr(sbox_in, 8)] for sbox_in in SM4.func_real_dec_SboxInID(self)], np.uint8)
    func_real_dec_SboxOutID = lambda self: SM4_SBOX[SM4.func_real_dec_SboxInID(self)]
    func_real_dec_SboxOutHW = lambda self: HW_ARRAY[SM4.func_real_dec_SboxOutID(self)]
    func_real_dec_SboxOutMBit = lambda self: np.array([[bit for bit in np.binary_repr(sbox_out, 8)] for sbox_out in SM4.func_real_dec_SboxOutID(self)], np.uint8)

    @staticmethod
    def _func_guess_enc_SboxInID(self):
        rkey_byte_need = int(self.target_byte / 4) * 4
        if self.rkey is None and rkey_byte_need > 0:
            raise RuntimeError(f'缺少轮密钥，请提供至少{rkey_byte_need}字节的轮密钥！')
        if rkey_byte_need > self.rkey.shape[1]:
            raise RuntimeError(f'轮密钥字节不足，请提供至少{rkey_byte_need}字节的轮密钥！')
        return sm4_sbox_in(self.plaintext, self.target_byte, self.rkey,True)

    @staticmethod
    def _func_get_guess_b0XORb02(self):
        guess_key = np.array(range(256), np.uint8)
        pt=self.plaintext[:,12]
        B0=SM4_SBOX[pt[:,None]^guess_key]
        B02=(B0^(B0<<2))&0xff
        xor=pt[:,None]^B02
        return xor

    func_guess_enc_SboxInID=lambda self:SM4._func_guess_enc_SboxInID(self)
    func_guess_enc_SboxInHW = lambda self: HW_ARRAY[SM4.func_guess_enc_SboxInID(self)]
    func_guess_enc_SboxInMBit=lambda self: np.array([[[bit for bit in np.binary_repr(sbox_in, 8)] for sbox_in in sbox_in_row] for sbox_in_row in SM4.func_guess_enc_SboxInID(self)], np.uint8)
    func_guess_enc_SboxOutID=lambda self: SM4_SBOX[SM4.func_guess_enc_SboxInID(self)]
    func_guess_enc_SboxOutHW = lambda self: HW_ARRAY[SM4.func_guess_enc_SboxOutID(self)]
    func_guess_enc_SboxOutMBit=lambda self: np.array([[[bit for bit in np.binary_repr(sbox_out, 8)] for sbox_out in sbox_out_row] for sbox_out_row in SM4.func_guess_enc_SboxOutID(self)], np.uint8)
    func_guess_enc_RoundXor = lambda self:SM4._func_get_guess_b0XORb02(self)
    func_guess_enc_RoundHD = lambda self:HW_ARRAY[SM4._func_get_guess_b0XORb02(self)]
    func_guess_enc_RoundMBit = lambda self:np.array([[[bit for bit in np.binary_repr(xor, 8)] for xor in xor_row] for xor_row in SM4._func_get_guess_b0XORb02(self)], np.uint8)

    @staticmethod
    def _func_guess_dec_SboxInID(self):
        rkey_byte_need = int(self.target_byte / 4) * 4
        if self.rkey is None and rkey_byte_need > 0:
            raise RuntimeError(f'缺少轮密钥，请提供至少{rkey_byte_need}字节的轮密钥！')
        if rkey_byte_need > self.rkey.shape[1]:
            raise RuntimeError(f'轮密钥字节不足，请提供至少{rkey_byte_need}字节的轮密钥！')
        return sm4_sbox_in(self.ciphertext, self.target_byte, self.rkey, True)

    func_guess_dec_SboxInID=lambda self:SM4._func_guess_dec_SboxInID(self)
    func_guess_dec_SboxInHW = lambda self: HW_ARRAY[SM4.func_guess_dec_SboxInID(self)]
    func_guess_dec_SboxInMBit = lambda self: np.array([[[bit for bit in np.binary_repr(sbox_in, 8)] for sbox_in in sbox_in_row] for sbox_in_row in SM4.func_guess_dec_SboxInID(self)], np.uint8)
    func_guess_dec_SboxOutID = lambda self: SM4_SBOX[SM4.func_guess_dec_SboxInID(self)]
    func_guess_dec_SboxOutHW = lambda self: HW_ARRAY[SM4.func_guess_dec_SboxOutID(self)]
    func_guess_dec_SboxOutMBit = lambda self: np.array([[[bit for bit in np.binary_repr(sbox_out, 8)] for sbox_out in sbox_out_row] for sbox_out_row in SM4.func_guess_dec_SboxOutID(self)], np.uint8)

    @staticmethod
    def _func_get_dec_rkey(self):
        rkey=np.array([sm4_key_expansion(mk,32) for mk in self.mkey],dtype=np.uint8)
        return np.hstack((rkey[:,-4:],rkey[:,-8:-4],rkey[:,-12:-8],rkey[:,-16:-12]))
    func_get_enc_rkey = lambda self: np.array([sm4_key_expansion(mk, 4) for mk in self.mkey], dtype=np.uint8)
    func_get_dec_rkey=lambda self:SM4._func_get_dec_rkey(self)

    _real_midvalue = {
        LEAKAGE_MODEL['SM4'][0]: func_real_enc_SboxInID,
        LEAKAGE_MODEL['SM4'][1]: func_real_enc_SboxInHW,
        LEAKAGE_MODEL['SM4'][2]: func_real_enc_SboxInMBit,
        LEAKAGE_MODEL['SM4'][3]: func_real_enc_SboxOutID,
        LEAKAGE_MODEL['SM4'][4]: func_real_enc_SboxOutHW,
        LEAKAGE_MODEL['SM4'][5]: func_real_enc_SboxOutMBit,
        LEAKAGE_MODEL['SM4'][6]: func_real_dec_SboxInID,
        LEAKAGE_MODEL['SM4'][7]: func_real_dec_SboxInHW,
        LEAKAGE_MODEL['SM4'][8]: func_real_dec_SboxInMBit,
        LEAKAGE_MODEL['SM4'][9]: func_real_dec_SboxOutID,
        LEAKAGE_MODEL['SM4'][10]: func_real_dec_SboxOutHW,
        LEAKAGE_MODEL['SM4'][11]: func_real_dec_SboxOutMBit,
        LEAKAGE_MODEL['SM4'][12]: func_real_enc_RoundXor,
        LEAKAGE_MODEL['SM4'][13]: func_real_enc_RoundHD,
        LEAKAGE_MODEL['SM4'][14]: func_real_enc_RoundMBit
    }

    _guess_midvalue = {
        LEAKAGE_MODEL['SM4'][0]: func_guess_enc_SboxInID,
        LEAKAGE_MODEL['SM4'][1]: func_guess_enc_SboxInHW,
        LEAKAGE_MODEL['SM4'][2]: func_guess_enc_SboxInMBit,
        LEAKAGE_MODEL['SM4'][3]: func_guess_enc_SboxOutID,
        LEAKAGE_MODEL['SM4'][4]: func_guess_enc_SboxOutHW,
        LEAKAGE_MODEL['SM4'][5]: func_guess_enc_SboxOutMBit,
        LEAKAGE_MODEL['SM4'][6]: func_guess_dec_SboxInID,
        LEAKAGE_MODEL['SM4'][7]: func_guess_dec_SboxInHW,
        LEAKAGE_MODEL['SM4'][8]: func_guess_dec_SboxInMBit,
        LEAKAGE_MODEL['SM4'][9]: func_guess_dec_SboxOutID,
        LEAKAGE_MODEL['SM4'][10]: func_guess_dec_SboxOutHW,
        LEAKAGE_MODEL['SM4'][11]: func_guess_dec_SboxOutMBit,
        LEAKAGE_MODEL['SM4'][12]: func_guess_enc_RoundXor,
        LEAKAGE_MODEL['SM4'][13]: func_guess_enc_RoundHD,
        LEAKAGE_MODEL['SM4'][14]: func_guess_enc_RoundMBit
    }

    _round_key = {
        LEAKAGE_MODEL['SM4'][0]: func_get_enc_rkey,
        LEAKAGE_MODEL['SM4'][1]: func_get_enc_rkey,
        LEAKAGE_MODEL['SM4'][2]: func_get_enc_rkey,
        LEAKAGE_MODEL['SM4'][3]: func_get_enc_rkey,
        LEAKAGE_MODEL['SM4'][4]: func_get_enc_rkey,
        LEAKAGE_MODEL['SM4'][5]: func_get_enc_rkey,
        LEAKAGE_MODEL['SM4'][6]: func_get_dec_rkey,
        LEAKAGE_MODEL['SM4'][7]: func_get_dec_rkey,
        LEAKAGE_MODEL['SM4'][8]: func_get_dec_rkey,
        LEAKAGE_MODEL['SM4'][9]: func_get_dec_rkey,
        LEAKAGE_MODEL['SM4'][10]: func_get_dec_rkey,
        LEAKAGE_MODEL['SM4'][11]: func_get_dec_rkey,
        LEAKAGE_MODEL['SM4'][12]: func_get_enc_rkey,
        LEAKAGE_MODEL['SM4'][13]: func_get_enc_rkey,
        LEAKAGE_MODEL['SM4'][14]: func_get_enc_rkey
    }

    @staticmethod
    def real_midvalue(leakage_model):
        return SM4._real_midvalue[leakage_model]

    @staticmethod
    def guess_midvalue(leakage_model):
        return SM4._guess_midvalue[leakage_model]

    @staticmethod
    def round_key(leakage_model):
        return SM4._round_key[leakage_model]

    @staticmethod
    def bytes_num(name):
        return {
            'plaintext_bytes': 16,
            'ciphertext_bytes': 16,
            'mkey_bytes': 16
        }[name]